package com.yupi.yuaiagent.chatmeory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 *
 */
public class FileBasedChatMemory implements ChatMemory {




    //
    private final String BASE_DIR;

    //
    private static final Kryo kryo = new Kryo();

    static {

        kryo.setRegistrationRequired(false);
        //设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    // 构造对象时,指定文件保存目录
    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }

 /*   *//**
     * 保存单条消息
     *
     * @param conversationId
     * @param message
     */
    @Override
    public void add(String conversationId, Message message) {
        List<Message> oneMessageList = new ArrayList<>();
        oneMessageList.add(message);
        saveConversation(conversationId, oneMessageList);
      //  saveConversation(conversationId, List.of(message));
    }


    /**
     * 根据 会话id, 添加多条数据
     *
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        //从文件中获取所有消息
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        //添加该消息列表到--消息列表中
        conversationMessages.addAll(messages);
        //将该会话id,对应的消息列表保存到该文件中
        saveConversation(conversationId, conversationMessages);


    }

    /**
     * 获取最后 --lastN 条数据
     *
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        //根据会话id,获取消息列表
        List<Message> messageList = getOrCreateConversation(conversationId);
        //跳过 --总条数 -当前的lastN
        return messageList.stream()
                .skip(Math.max(0, messageList.size() - lastN))
                .toList();

    }

    /**
     * 根据会话id。删除某个会话对应的文件
     *
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {



        //根据会话id,获取该文件
        File file = getConversationFile(conversationId);
        //文件存在就删除该文件
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 根据会话id, 获取消息列表
     *
     * @param conversationId
     * @return
     * @throws FileNotFoundException
     */
    private List<Message> getOrCreateConversation(String conversationId) {

        //获取该会话id的 文件
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (file.exists()) {
            //获取该kryo文件的信息。通过kryo方法解析成Messages 就可以了。
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    /**
     * 根据会话id. 保存消息
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        //根据会话id, 获取该文件
        File file = getConversationFile(conversationId);

        //使用 kryo 根据会话id.写入该文件最近的信息
        try (Output output = new Output(new FileOutputStream(file))) {

            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    /**
     * 每个会话单独文件保存
     *
     * @param conversationId
     * @return
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }
}
